DeepSpeed I/O, Offload, and Asynchrony

导言

“异步”只说明调用可以提前返回，不说明后台工作一定能被隐藏。DeepNVMe、Ulysses-Offload、ZenFlow 和 DataStates 分别搬运参数/状态、Attention 工作集、优化器更新与 checkpoint；判断它们是否有效，必须同时检查前台关键路径和后台队列是否稳定。


自绘示意图：四种机制搬运不同对象，并使用不同的完成语义。

统一判断：流水是否达到稳态

设每一步产生后台工作量 (D)，后台有效带宽为 (B)，相邻训练步间隔为 (T)。异步队列不发散至少要求

[
\frac{D}{B}\le T.
]

即使单步没有等待，只要 (D/B>T)，host cache、pinned buffer 或 NVMe 队列仍会持续积压，最终在某个 barrier 重新暴露。

DeepNVMe：先把 I/O 做成可并行的数据面

DeepNVMe 提供异步 NVMe 读写、pinned buffer、对齐和并行提交能力，是 ZeRO-Infinity 等 offload 路径的 I/O 基础设施。[^deepnvme]

机制

buffers = allocate_pinned_aligned_buffers(queue_depth)
for chunk_id in needed_chunks:
    slot = buffers.acquire()
    submit_async_read(file, offset(chunk_id), slot)

for chunk_id in compute_order:
    wait_read(chunk_id)
    async_copy_h2d(buffer(chunk_id))
    launch_gpu_compute(chunk_id)
    submit_async_write_if_dirty(chunk_id)

性能来自并行 outstanding request，而不是 Python async 语法。块太小会被 syscall/设备延迟主导，块太大又会降低预取灵活性；未锁页内存还会在 H2D/D2H 前增加隐式 staging。

边界

先用官方 ds_io benchmark 测设备在目标 block size、queue depth 和 parallel request 下的上限，再看训练是否接近该上限。共享文件系统、云盘或加密层可能使独立 NVMe 结论失效。

Ulysses-Offload：限制长序列 Attention 的活动工作集

Ulysses 把 [B,S/P,H,D] 通过 All-to-All 变成 [B,S,H/P,D]。超长序列下，即使每 Rank 只拥有部分 heads，完整序列的 Q/K/V、反向状态和 FFN/logits 仍会形成峰值。Ulysses-Offload 即 FPDT，把全局序列切为 chunk，并在 CPU 与 GPU 间双缓冲。[^fpdt-tutorial][^fpdt-paper]

for chunk in sequence_chunks:
    prefetch_next_chunk_to_gpu()
    qkv_local = project(local_slice(chunk))       # length C/P
    qkv_heads = ulysses_all_to_all(qkv_local)     # global length C, heads/P
    partial_out, lse = flash_attention(qkv_heads)
    merged_out = online_merge(merged_out, partial_out, lse)
    offload_backward_state_to_pinned_cpu()

固定 chunk (C) 只限制 chunk-local workspace：

[
M_{\mathrm{GPU}}\approx M_{\mathrm{model}}+
k_0\frac{BSM}{P}+k_1\frac{BCM}{P}.
]

总峰值仍含随 (S/P) 增长的本地 hidden-state 基座和模型/allocator floor；CPU 保存状态也会随层数和 (S/P) 增长。因此“任意长度”应理解为活动工作集受控，而不是容量与时间无限。

论文报告同硬件最高 16× 更长序列、4 GPU 训练 2M tokens、超过 55% MFU，均绑定其模型、硬件、chunk 与重计算口径。当前实现还受 Ulysses head 整除、序列/chunk 整除、FlashAttention 和自定义 autograd 边界约束。

{ width=92% }

FPDT Figure 4：红色块表示 pinned CPU 与 GPU 间移动的 Attention 状态；它解释方法，不代表所有 PCIe/NUMA 拓扑都能完全隐藏传输。

{ width=82% }

FPDT Figure 12(a)：2.7B、4 GPU、256K 序列的 chunk-size 消融。64K 是该实验的折中点，不是通用最优值。

ZenFlow：让 CPU optimizer 不必每步全量阻塞

ZeRO-Offload 可能把 GPU OOM 变成 CPU optimizer stall。ZenFlow 通过选择性梯度更新和异步 CPU optimizer，使一部分更新在 GPU 继续训练时推进。[^zenflow]

importance = score_gradient_partitions(local_gradients)
selected = topk_or_threshold(importance, config)
enqueue_cpu_optimizer(selected)
reuse_or_accumulate_unselected_partitions()
before_parameter_is_needed:
    wait_and_publish_required_update()

这里存在真实的数值近似：未选中的梯度/参数不会按普通 Adam 的逐步节奏更新。topk、更新间隔和异步陈旧度都可能影响收敛。因此验收必须包含：

• 相同 token/batch 的 loss 与下游质量；
• CPU optimizer 队列深度和 GPU 等待；
• 更新覆盖率、最大陈旧步数；
• 关闭 ZenFlow 的等配置对照。

DataStates：利用状态不可变窗口做 checkpoint

同步 checkpoint 往往执行 GPU → CPU → storage 后才恢复训练。DataStates-LLM 注意到某些模型/优化器状态在前向和反向窗口内不会被修改，可以把快照捕获与持久化解耦：先进入 host cache，再后台 flush。[^datastates-tutorial][^datastates-paper]

if checkpoint_due(step):
    reserve_host_snapshot(version=step)
    for state_shard in immutable_window:
        nonblocking_copy_to_host_cache(state_shard, version=step)
    publish_snapshot_metadata_when_complete(step)
    enqueue_storage_flush(step)

training_continues()
background_worker.flush_complete_snapshots_in_order()

正确性关键不是“文件最终写出来”，而是：

• 同一 checkpoint 的所有 shard 属于同一逻辑 step；
• host copy 完成前状态不会被复用或覆盖；
• 崩溃恢复只选择 metadata 已发布的完整版本；
• 后台积压不会耗尽 pinned host cache。

论文与教程报告最高 48× checkpoint 加速和 2.2× 端到端加速。当前教程明确列出 CUDA/NVIDIA、主要验证 ZeRO-1 且无 offload、尚不支持 universal/elastic checkpoint 等限制。


自绘示意图：当前台产生率超过后台消费率时，异步队列最终仍会回压训练。

四者不要串错对象

方法	搬运/延迟的对象	完成语义	主要风险
DeepNVMe	参数、优化器或任意 I/O chunk	read/write request 完成	对齐、queue depth、设备/FS 带宽
Ulysses-Offload	Q/K/V、反向状态、序列 tile	对应 chunk 可用于 forward/backward	PCIe 隐藏失败、CPU 容量、head/chunk 约束
ZenFlow	optimizer 更新	参数再次被消费前可见	近似与陈旧度影响收敛
DataStates	checkpoint shard	完整版本 metadata 原子发布	cache 积压、版本一致性、恢复兼容性

结论

先画出对象生命周期，再讨论异步。DeepNVMe 是数据面，Ulysses-Offload 是工作集流水，ZenFlow 是近似 optimizer 调度，DataStates 是持久化流水。每种方案都要同时通过“前台少等了多少”和“后台队列是否稳定”两道验收。

[^deepnvme]: DeepNVMe tutorial，教程文件最早提交于 2024-09-04。
[^fpdt-tutorial]: Ulysses-Offload tutorial，教程文件最早提交于 2024-12-05。
[^fpdt-paper]: FPDT paper。
[^zenflow]: ZenFlow tutorial，教程文件最早提交于 2025-08-15。
[^datastates-tutorial]: DataStates asynchronous checkpoint tutorial，教程文件最早提交于 2025-10-23。
[^datastates-paper]: DataStates-LLM paper。